#include <iostream>
#include <thread>
#include <chrono>
#include <cstring>
#include <vector>
#include <map>
#include <cmath>
// Linux网络编程头文件
#include <sys/socket.h>
#include <arpa/inet.h>
#include <unistd.h>
// ROS相关头文件
#include <ros/ros.h>

// 机器人状态结构体
struct RobotState {
    // 基础状态
    int robot_mode;          // 机器人模式
    int gait_type;          // 步态类型
    int battery_level;      // 电池电量
    int battery_current;    // 电池电流
    int battery_temperature;// 电池温度
    int battery_voltage;    // 电池电压

    // 位置和姿态
    float position[3];      // x, y, z位置
    float velocity[3];      // x, y, z速度
    float angular_velocity[3]; // roll, pitch, yaw角速度
    float attitude[3];      // roll, pitch, yaw姿态角

    // 足端状态
    int foot_force[4];      // 四足着地状态
    float foot_force_est[4];// 四足估计压力
    float foot_position[4][3]; // 四足位置
    float foot_speed[4][3];   // 四足速度

    // 电机状态
    int motor_temperature[12];  // 12个电机温度
    float motor_speed[12];     // 12个电机速度
    float motor_position[12];  // 12个电机位置
    float motor_current[12];   // 12个电机电流
    float motor_voltage[12];   // 12个电机电压

    // 错误状态
    unsigned int error_code;   // 错误代码
    unsigned int warning_code; // 警告代码

    // 派生状态
    bool is_power_on;         // 是否开机
    bool is_standing;         // 是否站立
    std::string current_mode; // 当前模式
    float speed;              // 当前速度
    bool is_moving;           // 是否在运动
};

// 心跳包发送类
class HeartBeat {
private:
    int udp_socket;
    struct sockaddr_in send_addr;
    bool running;
    std::thread heart_thread;

    void send_heartbeat() {
        while (running) {
            int data[3] = {0x21040001, 0, 0};
            sendto(udp_socket, (char*)data, sizeof(data), 0, 
                  (struct sockaddr*)&send_addr, sizeof(send_addr));
            std::this_thread::sleep_for(std::chrono::milliseconds(500));
        }
    }

public:
    HeartBeat(const char* ip = "192.168.1.120", int port = 43893) {
        // 创建UDP socket
        udp_socket = socket(AF_INET, SOCK_DGRAM, 0);
        if (udp_socket < 0) {
            ROS_ERROR("Failed to create heartbeat socket");
            return;
        }
        
        // 设置目标地址
        memset(&send_addr, 0, sizeof(send_addr));
        send_addr.sin_family = AF_INET;
        send_addr.sin_port = htons(port);
        send_addr.sin_addr.s_addr = inet_addr(ip);

        running = true;
        heart_thread = std::thread(&HeartBeat::send_heartbeat, this);
    }

    ~HeartBeat() {
        running = false;
        if (heart_thread.joinable()) {
            heart_thread.join();
        }
        close(udp_socket);
    }
};

// 状态监控类
class StateMonitor {
private:
    int udp_socket;
    struct sockaddr_in recv_addr;
    bool running;
    std::thread monitor_thread;
    RobotState state;

    void receive_state() {
        while (running) {
            char buffer[1024];
            struct sockaddr_in client_addr;
            socklen_t client_addr_len = sizeof(client_addr);
            
            int recv_len = recvfrom(udp_socket, buffer, sizeof(buffer), 0,
                                  (struct sockaddr*)&client_addr, &client_addr_len);
            
            if (recv_len > 0) {
                parse_state_data(buffer, recv_len);
            }
            
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }

    void parse_state_data(char* data, int length) {
        if (length < 4) return;

        int cmd_id = *(int*)data;
        if (cmd_id != 0x21040001) return;

        int offset = 4;

        // 解析基础状态
        memcpy(&state.robot_mode, data + offset, 24);
        offset += 24;

        // 解析位置和姿态
        memcpy(state.position, data + offset, 12);
        offset += 12;
        memcpy(state.velocity, data + offset, 12);
        offset += 12;
        memcpy(state.angular_velocity, data + offset, 12);
        offset += 12;
        memcpy(state.attitude, data + offset, 12);
        offset += 12;

        // 解析足端状态
        memcpy(state.foot_force, data + offset, 16);
        offset += 16;
        memcpy(state.foot_force_est, data + offset, 16);
        offset += 16;

        // 解析四足位置和速度
        for (int i = 0; i < 4; i++) {
            memcpy(state.foot_position[i], data + offset, 12);
            offset += 12;
        }
        for (int i = 0; i < 4; i++) {
            memcpy(state.foot_speed[i], data + offset, 12);
            offset += 12;
        }

        // 解析电机状态
        memcpy(state.motor_temperature, data + offset, 48);
        offset += 48;
        memcpy(state.motor_speed, data + offset, 48);
        offset += 48;
        memcpy(state.motor_position, data + offset, 48);
        offset += 48;
        memcpy(state.motor_current, data + offset, 48);
        offset += 48;
        memcpy(state.motor_voltage, data + offset, 48);
        offset += 48;

        // 解析错误和警告代码
        memcpy(&state.error_code, data + offset, 4);
        offset += 4;
        memcpy(&state.warning_code, data + offset, 4);

        // 更新派生状态
        update_derived_state();
    }

    void update_derived_state() {
        state.is_power_on = true;
        state.is_standing = (state.robot_mode == 1 || state.robot_mode == 2);
        
        switch (state.robot_mode) {
            case 1: state.current_mode = "spot"; break;
            case 2: state.current_mode = "move"; break;
            case 3: state.current_mode = "autonomous"; break;
            case 4: state.current_mode = "manual"; break;
            default: state.current_mode = "unknown";
        }

        state.speed = sqrt(state.velocity[0] * state.velocity[0] + 
                         state.velocity[1] * state.velocity[1]);
        state.is_moving = (state.speed > 0.01f);
    }

public:
    StateMonitor(const char* ip = "192.168.1.120", int port = 43894) {
        // 创建UDP socket
        udp_socket = socket(AF_INET, SOCK_DGRAM, 0);
        if (udp_socket < 0) {
            ROS_ERROR("Failed to create state monitor socket");
            return;
        }
        
        // 绑定地址
        memset(&recv_addr, 0, sizeof(recv_addr));
        recv_addr.sin_family = AF_INET;
        recv_addr.sin_port = htons(port);
        recv_addr.sin_addr.s_addr = inet_addr(ip);
        
        if (bind(udp_socket, (struct sockaddr*)&recv_addr, sizeof(recv_addr)) < 0) {
            ROS_ERROR("Failed to bind state monitor socket");
            close(udp_socket);
            return;
        }

        running = true;
        monitor_thread = std::thread(&StateMonitor::receive_state, this);
    }

    ~StateMonitor() {
        running = false;
        if (monitor_thread.joinable()) {
            monitor_thread.join();
        }
        close(udp_socket);
    }

    const RobotState& get_state() const { return state; }

    // 状态检查方法
    std::pair<bool, std::string> check_power() const {
        if (!state.is_power_on) {
            return {false, "机器狗未开机"};
        }
        return {true, "电源正常"};
    }

    std::pair<bool, std::string> check_battery(int min_level = 10) const {
        if (state.battery_level < min_level) {
            return {false, "电池电量过低"};
        }
        return {true, "电池电量正常"};
    }

    std::pair<bool, std::string> check_standing() const {
        if (!state.is_standing) {
            return {false, "机器狗必须处于站立状态"};
        }
        return {true, "站立状态正常"};
    }

    std::pair<bool, std::string> check_mode(const std::string& required_mode) const {
        if (state.current_mode != required_mode) {
            return {false, "模式不正确"};
        }
        return {true, "模式正常"};
    }

    std::pair<bool, std::string> check_posture(float max_angle = 30.0f) const {
        if (std::abs(state.attitude[0]) > max_angle || 
            std::abs(state.attitude[1]) > max_angle) {
            return {false, "姿态角度过大"};
        }
        return {true, "姿态正常"};
    }
};

// 机器人控制类
class RobotController {
private:
    int udp_socket;
    struct sockaddr_in send_addr;
    HeartBeat heartbeat;
    StateMonitor state_monitor;
    ros::NodeHandle nh;  // ROS节点句柄

    bool send_command(int code, int value, int type) {
        auto [power_ok, msg] = state_monitor.check_power();
        if (!power_ok) {
            ROS_ERROR_STREAM("错误：" << msg);
            return false;
        }

        int data[3] = {code, value, type};
        int sent = sendto(udp_socket, (char*)data, sizeof(data), 0,
                         (struct sockaddr*)&send_addr, sizeof(send_addr));
        return sent > 0;
    }

public:
    RobotController(const char* ip = "192.168.1.120", int port = 43893)
        : heartbeat(ip, port), state_monitor(ip, port + 1) {
        // 创建UDP socket
        udp_socket = socket(AF_INET, SOCK_DGRAM, 0);
        if (udp_socket < 0) {
            ROS_ERROR("Failed to create control socket");
            return;
        }
        
        // 设置目标地址
        memset(&send_addr, 0, sizeof(send_addr));
        send_addr.sin_family = AF_INET;
        send_addr.sin_port = htons(port);
        send_addr.sin_addr.s_addr = inet_addr(ip);
    }

    ~RobotController() {
        close(udp_socket);
    }

    // 基础控制命令
    bool stand_up_or_lie_down() {
        if (send_command(0x21010202, 0, 0)) {
            std::cout << "执行站立/趴下切换" << std::endl;
            return true;
        }
        return false;
    }

    bool soft_emergency_stop() {
        if (send_command(0x21020C0E, 0, 0)) {
            std::cout << "执行软急停" << std::endl;
            return true;
        }
        return false;
    }

    bool zero_position() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010C05, 0, 0)) {
            std::cout << "执行回零位" << std::endl;
            return true;
        }
        return false;
    }

    // 姿态调整命令（原地模式）
    bool adjust_roll(int value) {
        auto [stand_ok, stand_msg] = state_monitor.check_standing();
        auto [mode_ok, mode_msg] = state_monitor.check_mode("spot");
        if (!stand_ok || !mode_ok) {
            std::cout << "错误：" << (!stand_ok ? stand_msg : mode_msg) << std::endl;
            return false;
        }
        if (send_command(0x21010131, value, 0)) {
            std::cout << "调整横滚角，值：" << value << std::endl;
            return true;
        }
        return false;
    }

    bool adjust_pitch(int value) {
        auto [stand_ok, stand_msg] = state_monitor.check_standing();
        auto [mode_ok, mode_msg] = state_monitor.check_mode("spot");
        if (!stand_ok || !mode_ok) {
            std::cout << "错误：" << (!stand_ok ? stand_msg : mode_msg) << std::endl;
            return false;
        }
        if (send_command(0x21010130, value, 0)) {
            std::cout << "调整俯仰角，值：" << value << std::endl;
            return true;
        }
        return false;
    }

    bool adjust_height(int value) {
        auto [stand_ok, stand_msg] = state_monitor.check_standing();
        auto [mode_ok, mode_msg] = state_monitor.check_mode("spot");
        if (!stand_ok || !mode_ok) {
            std::cout << "错误：" << (!stand_ok ? stand_msg : mode_msg) << std::endl;
            return false;
        }
        if (send_command(0x21010102, value, 0)) {
            std::cout << "调整身体高度，值：" << value << std::endl;
            return true;
        }
        return false;
    }

    bool adjust_yaw(int value) {
        auto [stand_ok, stand_msg] = state_monitor.check_standing();
        auto [mode_ok, mode_msg] = state_monitor.check_mode("spot");
        if (!stand_ok || !mode_ok) {
            std::cout << "错误：" << (!stand_ok ? stand_msg : mode_msg) << std::endl;
            return false;
        }
        if (send_command(0x21010135, value, 0)) {
            std::cout << "调整偏航角，值：" << value << std::endl;
            return true;
        }
        return false;
    }

    // 移动控制命令
    bool move_side(int value) {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010131, value, 0)) {
            std::cout << "左右平移，速度值：" << value << std::endl;
            return true;
        }
        return false;
    }

    bool move_forward(int value) {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010130, value, 0)) {
            std::cout << "前后移动，速度值：" << value << std::endl;
            return true;
        }
        return false;
    }

    bool turn(int value) {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010135, value, 0)) {
            std::cout << "转向，速度值：" << value << std::endl;
            return true;
        }
        return false;
    }

    // 模式切换命令
    bool set_spot_mode() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010D05, 0, 0)) {
            std::cout << "切换到原地模式" << std::endl;
            return true;
        }
        return false;
    }

    bool set_move_mode() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010D06, 0, 0)) {
            std::cout << "切换到移动模式" << std::endl;
            return true;
        }
        return false;
    }

    // 步态控制命令
    bool set_flat_low_speed_gait() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010300, 0, 0)) {
            std::cout << "切换到平地低速步态" << std::endl;
            return true;
        }
        return false;
    }

    bool set_flat_medium_speed_gait() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010307, 0, 0)) {
            std::cout << "切换到平地中速步态" << std::endl;
            return true;
        }
        return false;
    }

    bool set_flat_high_speed_gait() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010303, 0, 0)) {
            std::cout << "切换到平地高速步态" << std::endl;
            return true;
        }
        return false;
    }

    // 特殊动作命令
    bool twist_body() {
        auto [ok, msg] = state_monitor.check_standing();
        if (!ok) {
            std::cout << "错误：" << msg << std::endl;
            return false;
        }
        if (send_command(0x21010204, 0, 0)) {
            std::cout << "执行扭身体动作" << std::endl;
            return true;
        }
        return false;
    }

    bool backflip() {
        auto [stand_ok, stand_msg] = state_monitor.check_standing();
        auto [batt_ok, batt_msg] = state_monitor.check_battery(50);
        if (!stand_ok || !batt_ok) {
            std::cout << "错误：" << (!stand_ok ? stand_msg : batt_msg) << std::endl;
            return false;
        }
        if (send_command(0x21010502, 0, 0)) {
            std::cout << "执行后空翻动作" << std::endl;
            return true;
        }
        return false;
    }

    // AI功能控制
    bool close_all_ai_options() {
        if (send_command(0x21012109, 0x00, 0)) {
            std::cout << "关闭所有AI功能" << std::endl;
            return true;
        }
        return false;
    }

    bool enable_obstacle_stop() {
        if (send_command(0x21012109, 0x20, 0)) {
            std::cout << "开启停障功能" << std::endl;
            return true;
        }
        return false;
    }

    bool enable_follow() {
        if (send_command(0x21012109, 0xC0, 0)) {
            std::cout << "开启跟随功能" << std::endl;
            return true;
        }
        return false;
    }
};

int main(int argc, char** argv) {
    // 初始化ROS节点
    ros::init(argc, argv, "dog_motion_controller");
    
    RobotController robot;
    ros::Duration(1.0).sleep();  // 使用ROS的延时函数
    
    robot.stand_up_or_lie_down();
    ros::Duration(2.0).sleep();
    
    robot.move_forward(10000);
    ros::Duration(2.0).sleep();
    
    // 保持节点运行
    ros::spin();
    
    return 0;
}
